/*
   Copyright 2011 Brook Smith

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/
package org.xempot.xml;

import java.util.ArrayDeque;
import java.util.Deque;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

import javax.xml.stream.XMLEventReader;
import javax.xml.stream.XMLStreamException;
import javax.xml.stream.events.Attribute;
import javax.xml.stream.events.XMLEvent;
import javax.xml.stream.util.EventReaderDelegate;

/**
 * Extends XmlEventReader to include context information, which is useful for
 * applying xpath style filters etc, within the bounds of a stream. <note>this
 * is not thread safe</note>
 * 
 * @author brook
 * 
 */
public class XmlContextEventReader extends EventReaderDelegate {

	private static final char SEPERATOR_CHAR = '/';

	private final Deque<String> stack = new ArrayDeque<String>();

	/**
	 * Map for maintaining the xpath locator
	 */
	private final Map<String, Map<String, Integer>> locators = new HashMap<String, Map<String, Integer>>();
	private String currentPath = "";
	/**
	 * Map for maintaining the hieracical attributes, i.e only the attributes of
	 * the elements direct ancestors are available
	 */
	private final Map<String, Map<String, Attribute>> attributes = new HashMap<String, Map<String, Attribute>>();

	private String flatPath = "";

	private boolean flatPathChanged = false;

	public XmlContextEventReader(XMLEventReader eventReader) {
		super(eventReader);
	}

	/**
	 * Traverses the tree back towards the root until the attribute is found if
	 * not found it will return null
	 * 
	 * @param key
	 * @return
	 */
	public XMLEvent getAttribute(String key) {
		if (key == null) {
			return null;
		}

		String fp = getFlatPath();
		while (fp.length() > 0) {
			if (attributes.get(fp).containsKey(key)) {
				return attributes.get(fp).get(key);
			}
			fp = trimPath(fp);
		}
		return null;
	}

	/**
	 * 
	 * @return an xpath style locator - without the index for each of an element
	 *         ie for /html/body/table[3]/tr[5] it will simply return
	 *         /html/body/table/tr
	 * 
	 */
	public String getFlatPath() {
		if (flatPathChanged) {
			String res = "";
			for (String element : stack) {
				res += SEPERATOR_CHAR + element;
			}
			flatPath = res;
			flatPathChanged = false;
		}
		return flatPath;
	}

	/**
	 * 
	 * @return xpath of the current element
	 */
	public String getXPathLocator() {
		return currentPath;
	}

	/**
	 * see {@linkplain XMLEventReader.nextEvent()}
	 */
	@Override
	public XMLEvent nextEvent() throws XMLStreamException {
		// maintain the flat locator path and the attributes
		XMLEvent event = super.nextEvent();
		String localPart = null;

		// why couldn't startElement and endElement inherit the same parent
		// element with a localpart?
		if (event.isStartElement()) {
			localPart = event.asStartElement().getName().getLocalPart();
			stack.push(localPart);
			flatPathChanged = true;

			String location = getFlatPath();
			// the iterator is a one shot iterator so transfer them into a map
			Iterator<Attribute> atts = event.asStartElement().getAttributes();
			Map<String, Attribute> attributeMap = new HashMap<String, Attribute>();
			while (atts.hasNext()) {
				Attribute attribute = atts.next();
				attributeMap.put(attribute.getName().getLocalPart(), attribute);
			}

			attributes.put(location, attributeMap);

		} else if (event.isEndElement()) {
			String location = getFlatPath();
			attributes.put(location, null);

			localPart = event.asEndElement().getName().getLocalPart();
			stack.pop();
			flatPathChanged = true;
		}

		// Maintain xpath locator
		if (event.isStartElement()) {
			Map<String, Integer> siblings = locators.get(currentPath);
			if (siblings == null) {
				siblings = new HashMap<String, Integer>();
				siblings.put(localPart, 1);
				locators.put(currentPath, siblings);
				currentPath += SEPERATOR_CHAR + localPart + "[1]";
			} else {
				// increment the localpath within parent
				Integer count = siblings.get(localPart);
				if (count == null) {
					count = 1;
				} else {
					count++;
				}
				siblings.put(localPart, count);
				currentPath += SEPERATOR_CHAR + localPart + "[" + count + "]";
			}
		} else if (event.isEndElement()) {
			// pop 1 off the locators
			locators.put(currentPath, null);
			// pop 1 off the current path
			currentPath = trimPath(currentPath);
		}

		return event;
	}

	private String trimPath(String path) {
		return path.substring(0, path.lastIndexOf(SEPERATOR_CHAR));
	}
}
